/*
 * Copyright 2021 TiKV Project Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package com.webank.tikv.operation;

import com.webank.tikv.PDClient;
import com.webank.tikv.exception.GrpcException;
import com.webank.tikv.exception.InternalException;
import com.webank.tikv.kvproto.Pdpb;
import com.webank.tikv.common.PDError;
import com.webank.tikv.util.BackOffer;
import io.grpc.StatusRuntimeException;

import java.util.function.Function;

import static com.webank.tikv.common.PDError.buildFromPdpbError;
import static com.webank.tikv.util.BackOffer.BoType.BoPDRPC;

public class PDErrorHandler<RespT> implements ErrorHandler<RespT> {

    public static final Function<Pdpb.GetRegionResponse, PDError> PD_RESPONSE_ERROR_FUNCTION =
        r -> r.getHeader().hasError() ? buildFromPdpbError(r.getHeader().getError())
            : r.getRegion().getId() == 0 ? PDError.RegionPeerNotElected.DEFAULT_INSTANCE : null;

    private final Function<RespT, PDError> getError;
    private final PDClient client;

    public PDErrorHandler(Function<RespT, PDError> errorExtractor, PDClient client) {
        this.getError = errorExtractor;
        this.client = client;
    }

    @Override
    public boolean handleResponseError(BackOffer backOffer, RespT resp) {
        if (resp == null) {
            return handleRequestError(backOffer, new GrpcException("PD Request Failed with unknown reason"));
        }
        PDError error = getError.apply(resp);
        if (error != null) {
            switch (error.getErrorType()) {
                case PD_ERROR:
                    boolean retry = backOffer.doBackOff(BoPDRPC, new GrpcException(error.toString()));
                    client.scheduleUpdateLeader();
                    return retry;
                case REGION_PEER_NOT_ELECTED:
                    return backOffer.doBackOff(BoPDRPC, new GrpcException(error.toString()));
                default:
                    throw new InternalException("Unknown error type encountered: " + error);
            }
        }
        return false;
    }

    @Override
    public boolean handleRequestError(BackOffer backOffer, Exception e) {
        // store id is not found
        if (e instanceof StatusRuntimeException && e.getMessage().contains("invalid store id")) {
            return false;
        }
        boolean retry = backOffer.doBackOff(BoPDRPC, e);
        client.scheduleUpdateLeader();
        return retry;
    }
}
